#include "../include/detector.hpp"
#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

//添加一些没用的注释以便测试，用来练习gitee功能
// 构造函数
JianChe::JianChe()
{ // 初始化 JianChe 对象
    // 相机内参矩阵
    camera_matrix = (Mat_<double>(3, 3) << 532.27294, 0., 327.3387,
                     0., 532.2387, 235.64438,
                     0, 0, 1);
    // 畸变系数
    dist_coeffs = (Mat_<double>(1, 5) << -0.044358, 0.124277, -0.002122, 0.002535, 0.000000);
}

// 识别装甲板并可视化
Mat JianChe::ShiBei(const Mat &inVideoFrame)
{
    Mat frame = inVideoFrame.clone();
    distances.clear();                               // 清空上一帧距离数据
    armor_centers.clear();                           // 清空上一帧中心点
    vector<Rect> armors_blue = findArmors(frame, 0); // 0:blue
    vector<Rect> armors_red = findArmors(frame, 1);  // 1:red

    // 合并结果
    vector<Rect> armors;
    armors.insert(armors.end(), armors_blue.begin(), armors_blue.end());
    armors.insert(armors.end(), armors_red.begin(), armors_red.end());

    for (const auto &armor : armors)
    {
        Scalar color;
        if (find(armors_blue.begin(), armors_blue.end(), armor) != armors_blue.end()) {
            color = ArmorColor(0);
        } else {
            color = ArmorColor(1);
        } // 区分颜色
        rectangle(frame, armor, color, 2);                              // 识别出来的色框

        string color_text = (color == Scalar(255, 0, 0)) ? "blue" : "red";
        putText(frame, "color = " + color_text, armor.tl(), FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 255, 255), 2); // 黄色文字
        // PnP测距逻辑
        // 1.定义装甲板实际3D坐标
        vector<Point3f> _3D_points = {
            {-67.5, -27.5, 0}, // 左上
            {67.5, -27.5, 0},  // 右上
            {67.5, 27.5, 0},   // 右下
            {-67.5, 27.5, 0}   // 左下
        };

        // 2.获取图像中装甲板四个角点
        vector<Point2f> image_points = {
            Point2f(armor.x, armor.y),
            Point2f(armor.x + armor.width, armor.y),
            Point2f(armor.x + armor.width, armor.y + armor.height),
            Point2f(armor.x, armor.y + armor.height)};

        // 检查内参是否有效
        if (camera_matrix.empty() || dist_coeffs.empty())
        {
            putText(frame, "内参未初始化", armor.tl() + Point(0, -50), FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 0, 255), 2);
            continue;
        }
        // 3.PnP解算
        Mat rvec, tvec;
        solvePnP(_3D_points, image_points, camera_matrix, dist_coeffs, rvec, tvec);

        // 4.计算距离（平移向量的模长）
        double distance = norm(tvec) / 1000.0;
        distances.push_back(distance);

        // 5.可视化距离
        string distance_text = "distance: " + to_string(distance) + "m";
        putText(frame, distance_text, Point(armor.x, armor.y - 30), FONT_HERSHEY_SIMPLEX, 1, Scalar(0, 255, 0), 2);

        // 6.记录中心点
        armor_centers.emplace_back(armor.x + armor.width / 2, armor.y + armor.height / 2);
    }
    return frame;
}

// 根据颜色区分装甲板
Scalar JianChe::ArmorColor(int color_flag)
{
    if (color_flag == 0)
        return Scalar(255, 0, 0); // 蓝色
    else
        return Scalar(0, 0, 255); // 红色
}

// 识别装甲板的具体函数
vector<Rect> JianChe::findArmors(const Mat &frame, int color_flag)
{
    vector<Rect> armors;
    Mat hsv;
    cvtColor(frame, hsv, COLOR_BGR2HSV); // 转换为 HSV 颜色空间

    // 定义蓝色和红色的HSV范围
    Mat mask;
    if (color_flag == 0)
    { // 蓝色
        Scalar lower_blue(90, 100, 100);
        Scalar upper_blue(128, 255, 255);
        inRange(hsv, lower_blue, upper_blue, mask);
    }
    else
    { // 红色
        Scalar lower_red1(0, 100, 100);
        Scalar upper_red1(10, 255, 255);
        Scalar lower_red2(156, 100, 100);
        Scalar upper_red2(180, 255, 255);

        Mat mask1, mask2;
        inRange(hsv, lower_red1, upper_red1, mask1);
        inRange(hsv, lower_red2, upper_red2, mask2);
        mask = mask1 | mask2; // [ADD] 红色有两个HSV区间
    }

    vector<vector<Point>> contours;
    findContours(mask, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE); // 将候选灯柱选出

    vector<Rect> possibleLightBars;
    for (const auto &contour : contours)
    {
        Rect rect = boundingRect(contour);

        int area = rect.width * rect.height;
        // 这里可以添加更多条件判断是否为装甲板，比如长宽比等
        if (rect.width > 10 && rect.height > 10 && area > 50)
        {
            possibleLightBars.push_back(rect);
        }
    }

    // 对可能的灯柱进行两两匹配
    for (size_t i = 0; i < possibleLightBars.size(); ++i)
    {
        for (size_t j = i + 1; j < possibleLightBars.size(); ++j)
        {
            if (isValidArmorPair(possibleLightBars[i], possibleLightBars[j]))
            {
                // 计算装甲板的包围框
                Rect armorRect = getArmorRect(possibleLightBars[i], possibleLightBars[j]);
                armors.push_back(armorRect);
            }
        }
    }

    return armors;
}

// 判断两个灯柱是否能组成有效的装甲板对
bool JianChe::isValidArmorPair(const Rect &lightBar1, const Rect &lightBar2)
{
    // 高度差限制
    double heightDiffRatio = std::abs(lightBar1.height - lightBar2.height) / static_cast<double>(std::max(lightBar1.height, lightBar2.height));
    if (heightDiffRatio > 0.2)
    {
        return false;
    }

    // 水平间距限制
    int centerX1 = lightBar1.x + lightBar1.width / 2;
    int centerX2 = lightBar2.x + lightBar2.width / 2;
    int horizontalDistance = std::abs(centerX1 - centerX2);
    double distanceRatio = horizontalDistance / static_cast<double>(std::max(lightBar1.height, lightBar2.height));
    if (distanceRatio < 0.5 || distanceRatio > 3)
    {
        return false;
    }

    // 垂直对齐限制
    int centerY1 = lightBar1.y + lightBar1.height / 2;
    int centerY2 = lightBar2.y + lightBar2.height / 2;
    int verticalDistance = std::abs(centerY1 - centerY2);
    double verticalRatio = verticalDistance / static_cast<double>(std::max(lightBar1.height, lightBar2.height));
    if (verticalRatio > 0.2)
    {
        return false;
    }

    return true;
}

// 计算由两个灯柱组成的装甲板的包围框
Rect JianChe::getArmorRect(const Rect &lightBar1, const Rect &lightBar2)
{
    int minX = std::min(lightBar1.x, lightBar2.x);
    int minY = std::min(lightBar1.y, lightBar2.y);
    int maxX = std::max(lightBar1.x + lightBar1.width, lightBar2.x + lightBar2.width);
    int maxY = std::max(lightBar1.y + lightBar1.height, lightBar2.y + lightBar2.height);
    return Rect(minX, minY, maxX - minX, maxY - minY);
}

// 判断是否为装甲板
bool JianChe::isArmor(const Mat &contourROI)
{
    int area = contourROI.cols * contourROI.rows;
    double aspectRatio = static_cast<double>(contourROI.cols) / contourROI.rows;
    // 判断逻辑
    if (area > 100 && aspectRatio > 0.5 && aspectRatio < 2)
    {
        return true;
    }
    return false;
}